In [94]:
!pip install catalyst --upgrade
!pip install --user git+https://github.com/albu/albumentations@bdd6a4e
!pip install --user git+https://github.com/qubvel/segmentation_models.pytorch
In [96]:
!pip install git+https://github.com/qubvel/segmentation_models.pytorch
In [1]:
import os
import cv2
import collections
import time
import tqdm
from PIL import Image
from functools import partial
train_on_gpu = True
In [2]:
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
%matplotlib inline
In [90]:
from sklearn.model_selection import train_test_split
from sklearn.metrics import roc_auc_score
import torchvision
import torchvision.transforms as transforms
import torch
from torch.utils.data import TensorDataset, DataLoader, Dataset
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.optim import lr_scheduler
from torch.utils.data.sampler import SubsetRandomSampler
from torch.optim.lr_scheduler import StepLR, ReduceLROnPlateau, CosineAnnealingLR
import segmentation_models_pytorch as smp
model = smp.Unet()
In [68]:
from albumentations import Compose, RandomCrop, Normalize, HorizontalFlip, Resize
from albumentations.pytorch import ToTensor
import albumentations as albu
from albumentations import pytorch as AT
In [148]:
from catalyst.dl.utils import criterion
In [151]:
from catalyst.data import Augmentor
from catalyst.dl import utils
from catalyst.data.reader import ImageReader, ScalarReader, ReaderCompose, LambdaReader
from catalyst.dl.runner import SupervisedRunner
#from catalyst.contrib.models.segmentation import Unet
from catalyst.dl.callbacks import DiceCallback, EarlyStoppingCallback, InferCallback, CheckpointCallback
In [44]:
def get_img(x, folder: str='train_images'):
"""
Return image based on image name and folder.
"""
data_folder = f"{path}/{folder}"
image_path = os.path.join(data_folder, x)
img = cv2.imread(image_path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
return img
def rle_decode(mask_rle: str = '', shape: tuple = (1400, 2100)):
'''
Decode rle encoded mask.
:param mask_rle: run-length as string formatted (start length)
:param shape: (height, width) of array to return
Returns numpy array, 1 - mask, 0 - background
'''
s = mask_rle.split()
starts, lengths = [np.asarray(x, dtype=int) for x in (s[0:][::2], s[1:][::2])]
starts -= 1
ends = starts + lengths
img = np.zeros(shape[0] * shape[1], dtype=np.uint8)
for lo, hi in zip(starts, ends):
img[lo:hi] = 1
return img.reshape(shape, order='F')
def make_mask(df: pd.DataFrame, image_name: str='img.jpg', shape: tuple = (1400, 2100)):
"""
Create mask based on df, image name and shape.
"""
encoded_masks = df.loc[df['im_id'] == image_name, 'EncodedPixels']
masks = np.zeros((shape[0], shape[1], 4), dtype=np.float32)
for idx, label in enumerate(encoded_masks.values):
if label is not np.nan:
mask = rle_decode(label)
masks[:, :, idx] = mask
return masks
def to_tensor(x, **kwargs):
"""
Convert image or mask.
"""
return x.transpose(2, 0, 1).astype('float32')
def mask2rle(img):
'''
Convert mask to rle.
img: numpy array, 1 - mask, 0 - background
Returns run length as string formated
'''
pixels= img.T.flatten()
pixels = np.concatenate([[0], pixels, [0]])
runs = np.where(pixels[1:] != pixels[:-1])[0] + 1
runs[1::2] -= runs[::2]
return ' '.join(str(x) for x in runs)
def visualize(image, mask, original_image=None, original_mask=None):
"""
Plot image and masks.
If two pairs of images and masks are passes, show both.
"""
fontsize = 14
class_dict = {0: 'Fish', 1: 'Flower', 2: 'Gravel', 3: 'Sugar'}
if original_image is None and original_mask is None:
f, ax = plt.subplots(1, 5, figsize=(24, 24))
ax[0].imshow(image)
for i in range(4):
ax[i + 1].imshow(mask[:, :, i])
ax[i + 1].set_title(f'Mask {class_dict[i]}', fontsize=fontsize)
else:
f, ax = plt.subplots(2, 5, figsize=(24, 12))
ax[0, 0].imshow(original_image)
ax[0, 0].set_title('Original image', fontsize=fontsize)
for i in range(4):
ax[0, i + 1].imshow(original_mask[:, :, i])
ax[0, i + 1].set_title(f'Original mask {class_dict[i]}', fontsize=fontsize)
ax[1, 0].imshow(image)
ax[1, 0].set_title('Transformed image', fontsize=fontsize)
for i in range(4):
ax[1, i + 1].imshow(mask[:, :, i])
ax[1, i + 1].set_title(f'Transformed mask {class_dict[i]}', fontsize=fontsize)
def visualize_with_raw(image, mask, original_image=None, original_mask=None, raw_image=None, raw_mask=None):
"""
Plot image and masks.
If two pairs of images and masks are passes, show both.
"""
fontsize = 14
class_dict = {0: 'Fish', 1: 'Flower', 2: 'Gravel', 3: 'Sugar'}
f, ax = plt.subplots(3, 5, figsize=(24, 12))
ax[0, 0].imshow(original_image)
ax[0, 0].set_title('Original image', fontsize=fontsize)
for i in range(4):
ax[0, i + 1].imshow(original_mask[:, :, i])
ax[0, i + 1].set_title(f'Original mask {class_dict[i]}', fontsize=fontsize)
ax[1, 0].imshow(raw_image)
ax[1, 0].set_title('Original image', fontsize=fontsize)
for i in range(4):
ax[1, i + 1].imshow(raw_mask[:, :, i])
ax[1, i + 1].set_title(f'Raw predicted mask {class_dict[i]}', fontsize=fontsize)
ax[2, 0].imshow(image)
ax[2, 0].set_title('Transformed image', fontsize=fontsize)
for i in range(4):
ax[2, i + 1].imshow(mask[:, :, i])
ax[2, i + 1].set_title(f'Predicted mask with processing {class_dict[i]}', fontsize=fontsize)
def plot_with_augmentation(image, mask, augment):
"""
Wrapper for `visualize` function.
"""
augmented = augment(image=image, mask=mask)
image_flipped = augmented['image']
mask_flipped = augmented['mask']
visualize(image_flipped, mask_flipped, original_image=image, original_mask=mask)
sigmoid = lambda x: 1 / (1 + np.exp(-x))
def post_process(probability, threshold, min_size):
"""
Post processing of each predicted mask, components with lesser number of pixels
than `min_size` are ignored
"""
# don't remember where I saw it
mask = cv2.threshold(probability, threshold, 1, cv2.THRESH_BINARY)[1]
num_component, component = cv2.connectedComponents(mask.astype(np.uint8))
predictions = np.zeros((350, 525), np.float32)
num = 0
for c in range(1, num_component):
p = (component == c)
if p.sum() > min_size:
predictions[p] = 1
num += 1
return predictions, num
def get_training_augmentation():
train_transform = [
albu.HorizontalFlip(p=0.5),
albu.ShiftScaleRotate(scale_limit=0.5, rotate_limit=0, shift_limit=0.1, p=0.5, border_mode=0),
albu.GridDistortion(p=0.5),
albu.OpticalDistortion(p=0.5, distort_limit=2, shift_limit=0.5),
albu.Resize(320, 640)
]
return albu.Compose(train_transform)
def get_validation_augmentation():
"""Add paddings to make image shape divisible by 32"""
test_transform = [
albu.Resize(320, 640)
]
return albu.Compose(test_transform)
def get_preprocessing(preprocessing_fn):
"""Construct preprocessing transform
Args:
preprocessing_fn (callbale): data normalization function
(can be specific for each pretrained neural network)
Return:
transform: albumentations.Compose
"""
_transform = [
albu.Lambda(image=preprocessing_fn),
albu.Lambda(image=to_tensor, mask=to_tensor),
]
return albu.Compose(_transform)
def dice(img1, img2):
img1 = np.asarray(img1).astype(np.bool)
img2 = np.asarray(img2).astype(np.bool)
intersection = np.logical_and(img1, img2)
return 2. * intersection.sum() / (img1.sum() + img2.sum())
In [45]:
path = '../cosmology'
os.listdir(path)
Out[45]:
In [46]:
train = pd.read_csv(f'{path}/train.csv')
sub = pd.read_csv(f'{path}/sample_submission.csv')
In [47]:
train.head()
Out[47]:
In [48]:
n_train = len(os.listdir(f'{path}/train_images'))
n_test = len(os.listdir(f'{path}/test_images'))
print(f'There are {n_train} images in train dataset')
print(f'There are {n_test} images in test dataset')
In [49]:
train['Image_Label'].apply(lambda x: x.split('_')[1]).value_counts()
Out[49]:
In [50]:
train.loc[train['EncodedPixels'].isnull() == False, 'Image_Label'].apply(lambda x: x.split('_')[1]).value_counts()
Out[50]:
In [51]:
train.loc[train['EncodedPixels'].isnull() == False, 'Image_Label'].apply(lambda x: x.split('_')[0]).value_counts().value_counts()
Out[51]:
In [52]:
train['label'] = train['Image_Label'].apply(lambda x: x.split('_')[1])
train['im_id'] = train['Image_Label'].apply(lambda x: x.split('_')[0])
sub['label'] = sub['Image_Label'].apply(lambda x: x.split('_')[1])
sub['im_id'] = sub['Image_Label'].apply(lambda x: x.split('_')[0])
In [53]:
fig = plt.figure(figsize=(25, 16))
for j, im_id in enumerate(np.random.choice(train['im_id'].unique(), 4)):
for i, (idx, row) in enumerate(train.loc[train['im_id'] == im_id].iterrows()):
ax = fig.add_subplot(5, 4, j * 4 + i + 1, xticks=[], yticks=[])
im = Image.open(f"{path}/train_images/{row['Image_Label'].split('_')[0]}")
plt.imshow(im)
mask_rle = row['EncodedPixels']
try: # label might not be there!
mask = rle_decode(mask_rle)
except:
mask = np.zeros((1400, 2100))
plt.imshow(mask, alpha=0.5, cmap='gray')
ax.set_title(f"Image: {row['Image_Label'].split('_')[0]}. Label: {row['label']}")
In [54]:
id_mask_count = train.loc[train['EncodedPixels'].isnull() == False, 'Image_Label'].apply(lambda x: x.split('_')[0]).value_counts().\
reset_index().rename(columns={'index': 'img_id', 'Image_Label': 'count'})
train_ids, valid_ids = train_test_split(id_mask_count['img_id'].values, random_state=42, stratify=id_mask_count['count'], test_size=0.1)
test_ids = sub['Image_Label'].apply(lambda x: x.split('_')[0]).drop_duplicates().values
In [55]:
image_name = '8242ba0.jpg'
image = get_img(image_name)
mask = make_mask(train, image_name)
In [56]:
visualize(image, mask)
In [57]:
plot_with_augmentation(image, mask, albu.HorizontalFlip(p=1))
In [58]:
plot_with_augmentation(image, mask, albu.VerticalFlip(p=1))
In [59]:
plot_with_augmentation(image, mask, albu.RandomRotate90(p=1))
In [60]:
plot_with_augmentation(image, mask, albu.ElasticTransform(p=1, alpha=120, sigma=120 * 0.05, alpha_affine=120 * 0.03))
In [63]:
plot_with_augmentation(image, mask, albu.GridDistortion(p=0.1))
In [64]:
plot_with_augmentation(image, mask, albu.OpticalDistortion(p=1, distort_limit=2, shift_limit=0.5))
In [71]:
class CloudDataset(Dataset):
def __init__(self, df: pd.DataFrame = None, datatype: str = 'train', img_ids: np.array = None,
transforms = albu.Compose([albu.HorizontalFlip(),AT.ToTensor()]),
preprocessing=None):
self.df = df
if datatype != 'test':
self.data_folder = f"{path}/train_images"
else:
self.data_folder = f"{path}/test_images"
self.img_ids = img_ids
self.transforms = transforms
self.preprocessing = preprocessing
def __getitem__(self, idx):
image_name = self.img_ids[idx]
mask = make_mask(self.df, image_name)
image_path = os.path.join(self.data_folder, image_name)
img = cv2.imread(image_path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
augmented = self.transforms(image=img, mask=mask)
img = augmented['image']
mask = augmented['mask']
if self.preprocessing:
preprocessed = self.preprocessing(image=img, mask=mask)
img = preprocessed['image']
mask = preprocessed['mask']
return img, mask
def __len__(self):
return len(self.img_ids)
In [72]:
ENCODER = 'resnet50'
ENCODER_WEIGHTS = 'imagenet'
DEVICE = 'cuda'
ACTIVATION = None
model = smp.Unet(
encoder_name=ENCODER,
encoder_weights=ENCODER_WEIGHTS,
classes=4,
activation=ACTIVATION,
)
preprocessing_fn = smp.encoders.get_preprocessing_fn(ENCODER, ENCODER_WEIGHTS)
In [144]:
num_workers = 0
bs = 16
train_dataset = CloudDataset(df=train, datatype='train', img_ids=train_ids, transforms = get_training_augmentation(), preprocessing=get_preprocessing(preprocessing_fn))
valid_dataset = CloudDataset(df=train, datatype='valid', img_ids=valid_ids, transforms = get_validation_augmentation(), preprocessing=get_preprocessing(preprocessing_fn))
train_loader = DataLoader(train_dataset, batch_size=bs, shuffle=True, num_workers=num_workers)
valid_loader = DataLoader(valid_dataset, batch_size=bs, shuffle=False, num_workers=num_workers)
loaders = {
"train": train_loader,
"valid": valid_loader
}
In [145]:
In [158]:
num_epochs = 19
logdir = "./logs/segmentation"
# model, criterion, optimizer
optimizer = torch.optim.Adam([
{'params': model.decoder.parameters(), 'lr': 1e-2},
{'params': model.encoder.parameters(), 'lr': 1e-3},
])
scheduler = ReduceLROnPlateau(optimizer, factor=0.15, patience=2)
#criterion = smp.utils.losses.BCELoss(None)
criterion = nn.CrossEntropyLoss()
runner = SupervisedRunner()
In [159]:
runner = SupervisedRunner()
runner.train(
model=model,
criterion=criterion,
optimizer=optimizer,
scheduler=scheduler,
loaders=loaders,
callbacks=[DiceCallback(), EarlyStoppingCallback(patience=5, min_delta=0.001)],
logdir=logdir,
num_epochs=num_epochs,
verbose=True
)
In [88]:
utils.plot_metrics(
logdir=logdir,
# specify which metrics we want to plot
metrics=["loss", "dice", 'lr', '_base/lr']
)
In [160]:
encoded_pixels = []
loaders = {"infer": valid_loader}
runner.infer(
model=model,
loaders=loaders,
callbacks=[
CheckpointCallback(
resume=f"{logdir}/checkpoints/best.pth"),
InferCallback()
],
)
valid_masks = []
probabilities = np.zeros((2220, 350, 525))
for i, (batch, output) in enumerate(tqdm.tqdm(zip(
valid_dataset, runner.callbacks[0].predictions["logits"]))):
image, mask = batch
for m in mask:
if m.shape != (350, 525):
m = cv2.resize(m, dsize=(525, 350), interpolation=cv2.INTER_LINEAR)
valid_masks.append(m)
for j, probability in enumerate(output):
if probability.shape != (350, 525):
probability = cv2.resize(probability, dsize=(525, 350), interpolation=cv2.INTER_LINEAR)
probabilities[i * 4 + j, :, :] = probability
In [99]:
torch.cuda.is_available
Out[99]:
In [100]:
x.cuda()
In [ ]: